
import datetime

from helper import read_file_content, init_Ark_client

#
model = "ep-20250318155257-7qkch"
mode = "common_prefix"

def create_context(client, background_content):
    print("----- create context -----")
    response = client.context.create(
        model=model,
        mode=mode,
        messages=[
            {"role": "system", "content": background_content},
        ],
        ttl=datetime.timedelta(minutes=600),
    )
    print(response)
    return response

def non_stream_chat(client, context_id, question_content):
    print("----- chat round 1 (non-stream) -----")
    chat_response = client.context.completions.create(
        # 指定上下文ID
        context_id=context_id,
        # 指定模型
        model=model,
        # 设置消息列表，包含一个用户角色的消息
        messages=[
            {"role": "user", "content": question_content},
        ],
        # 设置为非流式
        stream=False,
    )
    print(chat_response.choices[0].message.content)
    print(chat_response.usage)

def stream_chat(client, context_id, question_content):
    print("----- chat round 2 (streaming) -----")
    # 调用client的context.completions.create方法进行流式聊天
    stream = client.context.completions.create(
        # 指定上下文ID
        context_id=context_id,
        # 指定模型
        model=model,
        # 设置消息列表，包含一个用户角色的消息
        messages=[
            {"role": "user", "content": question_content},
        ],
        # 设置流式选项，包含使用情况
        stream_options={
            "include_usage": True,
        },
        # 设置为流式
        stream=True,
    )
    # 遍历流式响应的每个块
    for chunk in stream:
        # 如果块中包含使用情况信息，则打印使用情况
        if chunk.usage:
            print(chunk.usage)
        # 如果块中不包含选择信息，则继续下一个块
        if not chunk.choices:
            continue
        print(chunk.choices[0].delta.content, end="")

# 如果当前脚本是作为主程序运行
def main():
    background_content = read_file_content("background.txt")
    question_content = read_file_content("question.txt")
    client = init_Ark_client()
    context_response = create_context(client, background_content)
    non_stream_chat(client, context_response.id, question_content)

if __name__ == "__main__":
    main()
